#!/usr/bin/env python3
"""
@Author: wujin_xueyk\n
@Date:   2022-01-01\n
@Info: this is a library for rotation transform\n
@Ref:\n
1. 《小型旋翼无人机姿态融合及导航方法》\n
2. https://docs.microsoft.com/en-us/windows/win32/api/directxmath/nf-directxmath-xmquaternionrotationmatrix\n
"""

import math
import numpy as np

##########################################################################
# Basic
##########################################################################


def R_inv(R):
    """
    Get the inverse R(DCM)

    Inputs:
        - R: two-demision array[3x3], Rotation matrix
    Outputs:
        R: two-demision array[3x3], inverse of input Rotation matrix
    """
    R_inv = [[R[0][0], R[1][0], R[2][0]],
             [R[0][1], R[1][1], R[2][1]]
             [R[0][2], R[1][2], R[2][2]]]
    return R_inv


def q_mult(qa, qb):
    """
    Calculate Quaternion multiplication

    Inputs:
        - qa: array[4x1], Quaternion
        - qb: array[4x1], Quaternion
    Outputs:
        q: array[4x1], inverse of input Quaternion
    """
    qa_m = [[qa[0], -qa[1], -qa[2], -qa[3]],
            [qa[1],  qa[0], -qa[3],  qa[2]],
            [qa[2],  qa[3],  qa[0], -qa[1]],
            [qa[3], -qa[2],  qa[1],  qa[0]]]

    qa_m = np.array(qa_m)
    qb = np.array(qb)
    q = np.dot(qa_m, qb)
    q = q.tolist()
    return q


def q_inv(q):
    """
    Calculate the inverse of Quaternion

    Inputs:
        - q: array[4x1], Quaternion
    Outputs:
        q: array[4x1], inverse of input Quaternion
    """
    q0 = q[0]
    q1 = q[1]
    q2 = q[2]
    q3 = q[3]

    q_inv = [q0, -q1, -q2, -q3]
    return q_inv


def v_cross(v):
    """
    Get the vector cross matrix
    """
    vx = v[0]
    vy = v[1]
    vz = v[2]
    v_m = [[0, -vz, vy],
           [vz, 0, -vx],
           [-vy, vx, 0]]

    return v_m
##########################################################################
# Convert euler, R(DCM), and Quaternion
##########################################################################


def euler2Rbe(euler):
    """
    Convert Euler angle to Rotation matrix from body frame to earth frame. \n
    1. Euler angle defined in NED frame.\n
    2. Take v_b is vector[3x1] in body frame, then v_e = Rbe*v_b.\n

    Inputs:
        - euler: a tuple, (roll, pitch, yaw), in radian unit.
    Outputs:
        Rbe: two-demision array[3x3], Rotation matrix from body frame to earth frame
    """
    roll, pitch, yaw = euler[0], euler[1], euler[2]
    Rbe = [[math.cos(pitch)*math.cos(yaw), math.sin(roll)*math.sin(pitch)*math.cos(yaw) - math.cos(roll)*math.sin(yaw), math.cos(roll)*math.sin(pitch)*math.cos(yaw) + math.sin(roll)*math.sin(yaw)],
           [math.cos(pitch)*math.sin(yaw), math.sin(roll)*math.sin(pitch)*math.sin(yaw) + math.cos(roll)
            * math.cos(yaw), math.cos(roll)*math.sin(pitch)*math.sin(yaw) - math.sin(roll)*math.cos(yaw)],
           [-math.sin(pitch),               math.sin(roll)*math.cos(pitch),                                              math.cos(roll)*math.cos(pitch)]]

    return Rbe


def euler2Reb(euler):
    """
    Convert Euler angle to Rotation matrix from earth frame to body frame. \n
    1. Euler angle defined in NED frame.\n
    2. Take v_e is vector[3x1] in earth frame, then v_b = Rbe*v_e.\n

    Inputs:
        - euler: a tuple, (roll, pitch, yaw), in radian unit.
    Outputs:
        Reb: two-demision array[3x3], Rotation matrix from earth frame to body frame
    """
    roll, pitch, yaw = euler[0], euler[1], euler[2]
    Reb = [[math.cos(pitch)*math.cos(yaw),                                               math.cos(pitch)*math.sin(yaw),                                              -math.sin(pitch)],
           [math.sin(roll)*math.sin(pitch)*math.cos(yaw) - math.cos(roll)*math.sin(yaw), math.sin(roll) *
            math.sin(pitch)*math.sin(yaw) + math.cos(roll)*math.cos(yaw), math.sin(roll)*math.cos(pitch)],
           [math.cos(roll)*math.sin(pitch)*math.cos(yaw) + math.sin(roll)*math.sin(yaw), math.cos(roll)*math.sin(pitch)*math.sin(yaw) - math.sin(roll)*math.cos(yaw), math.cos(roll)*math.cos(pitch)]]

    return Reb


def Rbe2euler(Rbe):
    """
    Convert Rotation matrix from body frame to earth frame to Euler angle. \n
    1. Euler angle defined in NED frame.\n
    2. Take v_b is vector[3x1] in body frame, then v_e = Rbe*v_b.\n

    Inputs:
        - Rbe: two-demision array[3x3], Rotation matrix from body frame to earth frame
    Outputs:
        euler: a tuple, (roll, pitch, yaw), in radian unit.
    """
    roll = math.atan2(Rbe[2][1], Rbe[2][2])
    pitch = -math.asin(Rbe[2][0])
    yaw = math.atan2(Rbe[1][0], Rbe[0][0])
    return (roll, pitch, yaw)


def Reb2euler(Reb):
    """
    Convert Rotation matrix from earth frame to body frame to Euler angle. \n
    1. Euler angle defined in NED frame.\n
    2. Take v_e is vector[3x1] in earth frame, then v_b = Rbe*v_e.\n

    Inputs:
        - Reb: two-demision array[3x3], Rotation matrix from earth frame to body frame
    Outputs:
        euler: a tuple, (roll, pitch, yaw), in radian unit.
    """
    roll = math.atan2(Reb[1][2], Reb[2][2])
    pitch = -math.asin(Reb[0][2])
    yaw = math.atan2(Reb[0][1], Reb[0][0])
    return (roll, pitch, yaw)


def q2R(q):
    """
    Convert Quaternion to Rotation matrix. \n
    1. Quaternion = [w, x, y, z], where w=cos(angle/2), x=x*sin(angle/2), y=y*sin(angle/2), z=z*sin(angle/2), and vector [x,y,z] is in body frame.\n

    Inputs:
        - q: array[4x1], Quaternion
    Outputs:
        R: two-demision array[3x3], Rotation matrix
    """
    q0 = q[0]   # w
    q1 = q[1]   # x
    q2 = q[2]   # y
    q3 = q[3]   # z

    R = [[q0*q0+q1*q1-q2*q2-q3*q3, 2*(q1*q2-q0*q3),         2*(q1*q3+q0*q2)],
         [2*(q1*q2+q0*q3),         q0*q0-q1*q1+q2*q2-q3*q3, 2*(q2*q3-q0*q1)],
         [2*(q1*q3-q0*q2),         2*(q2*q3+q0*q1),         q0*q0-q1*q1-q2*q2+q3*q3]]

    return R


def euler2qbe(euler):
    """
    Convert euler to Quaternion from body frame to earth frame.\n
    1. Euler angle defined in NED frame.\n
    2. Take v_b is vector[3x1] in body frame, then v_e = qbe x [0, v_b] x inv(qbe).\n

    Inputs:
        - euler: a tuple, (roll, pitch, yaw), in radian unit.
    Outputs:
        qbe: array[4x1], Quaternion
    """
    roll, pitch, yaw = euler[0], euler[1], euler[2]
    q0 = np.cos(roll/2)*np.cos(pitch/2)*np.cos(yaw/2) + \
        np.sin(roll/2)*np.sin(pitch/2)*np.sin(yaw/2)
    q1 = np.sin(roll/2)*np.cos(pitch/2)*np.cos(yaw/2) - \
        np.cos(roll/2)*np.sin(pitch/2)*np.sin(yaw/2)
    q2 = np.cos(roll/2)*np.sin(pitch/2)*np.cos(yaw/2) + \
        np.sin(roll/2)*np.cos(pitch/2)*np.sin(yaw/2)
    q3 = np.cos(roll/2)*np.cos(pitch/2)*np.sin(yaw/2) - \
        np.sin(roll/2)*np.sin(pitch/2)*np.cos(yaw/2)

    qbe = [q0, q1, q2, q3]
    return qbe


def euler2qeb(euler):
    """
    Convert euler to Quaternion from earth frame to body frame.\n
    1. Euler angle defined in NED frame.\n
    2. Take v_e is vector[3x1] in earth frame, then v_b = qeb x [0, v_e] x inv(qeb).\n

    Inputs:
        - euler: a tuple, (roll, pitch, yaw), in radian unit.
    Outputs:
        qeb: array[4x1], Quaternion
    """
    roll, pitch, yaw = euler[0], euler[1], euler[2]
    q0 = np.cos(roll/2)*np.cos(pitch/2)*np.cos(yaw/2) + \
        np.sin(roll/2)*np.sin(pitch/2)*np.sin(yaw/2)
    q1 = -np.sin(roll/2)*np.cos(pitch/2)*np.cos(yaw/2) + \
        np.cos(roll/2)*np.sin(pitch/2)*np.sin(yaw/2)
    q2 = -np.cos(roll/2)*np.sin(pitch/2)*np.cos(yaw/2) - \
        np.sin(roll/2)*np.cos(pitch/2)*np.sin(yaw/2)
    q3 = -np.cos(roll/2)*np.cos(pitch/2)*np.sin(yaw/2) + \
        np.sin(roll/2)*np.sin(pitch/2)*np.cos(yaw/2)

    qeb = [q0, q1, q2, q3]
    return qeb


def qbe2euler(qbe):
    """
    Convert Quaternion from body frame to earth frame to euler.\n
    1. Euler angle defined in NED frame.\n
    2. Take v_b is vector[3x1] in body frame, then v_e = qbe x [0, v_b] x inv(qbe).\n

    Inputs:
        - qbe: array[4x1], Quaternion
    Outputs:
        euler: a tuple, (roll, pitch, yaw), in radian unit.
    """
    Rbe = q2R(qbe)
    euler = Rbe2euler(Rbe)
    return euler


def qeb2euler(qeb):
    """
    Convert Quaternion from earth frame to body frame to euler.\n
    1. Euler angle defined in NED frame.\n
    2. Take v_e is vector[3x1] in earth frame, then v_b = qeb x [0, v_e] x inv(qeb).\n

    Inputs:
        - euler: a tuple, (roll, pitch, yaw), in radian unit.
    Outputs:
        qeb: array[4x1], Quaternion
    """
    Reb = q2R(qeb)
    euler = Reb2euler(Reb)
    return euler

##########################################################################
# Transfer/Mapping vector by R(DCM) or Quaternion
##########################################################################


def tf_with_q(q, v):
    """
    Transfer/Mapping the vector with Quaternion.
    """
    vx = v[0]
    vy = v[1]
    vz = v[2]
    v_q = [0, vx, vy, vz]
    tmp = q_mult(q, v_q)
    tmp = q_mult(tmp, q_inv(q))

    res = [tmp[1], tmp[2], tmp[3]]
    return res


def tf_with_R(R, v):
    """
    Transfer/Mapping the vector with Rotation matrix.
    """
    R_m = np.array(R)
    v_m = np.array(v)
    res = np.dot(R_m, v_m)

    res = res.tolist()
    return res

##########################################################################
# Transfer R(DCM) or Quaternion
##########################################################################


def Rbe_ENU2NED(Rbe1):
    """
    Get Rbe2 in NED frame base on the Rbe1 in ENU frame\n

    Inputs:
        - Rbe1: the DCM in ENU frame
    Outputs:
        Rbe2: the DCM in NED frame
    """
    # from body(NED) to body(ENU)
    R_b2_2_b1 = np.array([[1, 0, 0],
                          [0, -1, 0],
                          [0, 0, -1]])
    # from earth(NED) to earth(ENU)
    R_e2_2_e1 = np.array([[0, 1, 0],
                          [1, 0, 0],
                          [0, 0, -1]])
    # from earth(ENU) to earth(NED)
    R_e1_2_e2 = np.linalg.inv(R_e2_2_e1)

    Rbe1 = np.array(Rbe1)
    Rbe2 = np.dot(R_e1_2_e2, np.dot(Rbe1, R_b2_2_b1))
    Rbe2 = Rbe2.tolist()
    return Rbe2


def Rbe_NWU2NED(Rbe1):
    """
    Get R in NED frame base on the R in baselink(NWU) frame\n

    Inputs:
        - R1: the DCM in baselink frame
    Outputs:
        R2: the DCM in NED frame
    """
    # from body(NED) to body(NWU)
    R_b2_2_b1 = np.array([[1, 0, 0],
                          [0, -1, 0],
                          [0, 0, -1]])
    # from earth(NED) to earth(NWU)
    R_e2_2_e1 = np.array([[1, 0, 0],
                          [0, -1, 0],
                          [0, 0, -1]])
    # from earth(NWU) to earth(NED)
    R_e1_2_e2 = np.linalg.inv(R_e2_2_e1)

    #
    Rbe1 = np.array(Rbe1)
    Rbe2 = np.dot(R_e2_2_e1, np.dot(Rbe1, np.linalg.inv(R_b2_2_b1)))

    Rbe2 = Rbe2.tolist()
    return Rbe2


def qbe_ENU2NED(qbe1):
    """
    Get qbe in NED frame base on the q in ENU frame.\n

    Inputs:
        - qbe1: quaternion from body to earth in ENU frame
    Outputs:
        qbe2: quaternion from body to earth in NED frame
    """
    # quaternion from body2 to body1: equal to the rotation to let body1 frame coincide with body2 frame
    # step1: rotate 180deg around x-axis in body frame
    theta = np.pi
    q_b2_2_b1 = np.array([np.cos(theta/2), np.sin(theta/2), 0, 0])

    # quaternion from earth2 to earth1: equal to the rotation to let earth1 frame coincide with earth2 frame
    # step1: rotate 180deg around y-axis in ENU frame
    theta = np.pi
    q_1 = np.array([np.cos(theta/2), 0, np.sin(theta/2), 0])
    # step2: rotate 90deg around z-axis in ENU frame
    theta = np.pi/2
    q_2 = np.array([np.cos(theta/2), 0, 0, np.sin(theta/2)])
    q_e2_2_e1 = q_mult(q_1, q_2)
    q_e1_2_e2 = q_inv(q_e2_2_e1)

    q_ = q_mult(qbe1, q_b2_2_b1)
    qbe2 = q_mult(q_e1_2_e2, q_)
    return qbe2


def qbe_NED2ENU(qbe1):
    """
    Get qbe in ENU frame base on the q in NED frame.\n

    Inputs:
        - qbe1: quaternion from body to earth in NED frame
    Outputs:
        qbe2: quaternion from body to earth in ENU frame
    """
    # quaternion from body2 to body1: equal to the rotation to let body1 frame coincide with body2 frame
    # step1: rotate 180deg around x-axis in body frame
    theta = np.pi
    q_b2_2_b1 = np.array([np.cos(theta/2), np.sin(theta/2), 0, 0])

    # quaternion from earth2 to earth1: equal to the rotation to let earth1 frame coincide with earth2 frame
    # step1: rotate 180deg around y-axis in NED frame
    theta = np.pi
    q_1 = np.array([np.cos(theta/2), 0, np.sin(theta/2), 0])
    # step2: rotate 90deg around z-axis in NED frame
    theta = np.pi/2
    q_2 = np.array([np.cos(theta/2), 0, 0, np.sin(theta/2)])
    q_e2_2_e1 = q_mult(q_1, q_2)
    q_e1_2_e2 = q_inv(q_e2_2_e1)

    q_ = q_mult(qbe1, q_b2_2_b1)
    qbe2 = q_mult(q_e1_2_e2, q_)
    return qbe2


def qbe_NWU2NED(qbe1):
    """
    Get qbe in NED frame base on the qbe in baselink(NWU) frame.\n

    Inputs:
        - qbe1: quaternion from body to earth in baselink(NWU) frame
    Outputs:
        qbe2: quaternion from body to earth in NED frame
    """
    # quaternion from body2 to body1: equal to the rotation to let body1 frame coincide with body2 frame
    # step1: rotate 180deg around x-axis in body frame
    theta = np.pi
    q_b2_2_b1 = np.array([np.cos(theta/2), np.sin(theta/2), 0, 0])

    # quaternion from earth2 to earth1: equal to the rotation to let earth1 frame coincide with earth2 frame
    # step1: rotate 180deg around x-axis in ENU frame
    theta = np.pi
    q_e2_2_e1 = np.array([np.cos(theta/2), np.sin(theta/2), 0, 0])
    q_e1_2_e2 = q_inv(q_e2_2_e1)

    q_ = q_mult(qbe1, q_b2_2_b1)
    qbe2 = q_mult(q_e1_2_e2, q_)
    return qbe2


def test1():
    euler = (-1.0, -0.4, 0.0)
    # Rbe = euler2Rbe(euler)
    # euler = Rbe2euler(Rbe)
    ###
    Rbe = euler2Reb(euler)
    euler = Reb2euler(Rbe)

    print(euler)


def test2():
    euler = (1.0, -0.4, 0.0)
    # qbe = euler2qbe(euler)
    # euler = qbe2euler(qbe)
    ###
    qeb = euler2qeb(euler)
    euler = qeb2euler(qeb)

    print(euler)


def main():
    # test1()
    test2()


if __name__ == '__main__':
    main()
